{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp tslearner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TSLearners (TSClassifier, TSRegressor, TSForecaster)\n",
    "\n",
    "> New set of time series learners with a new sklearn-like API that simplifies the learner creation.\n",
    "\n",
    "The new APIs make it easier to create a classifier, a regressor or a forecaster. You will no longer need to remember what arguments go to the dataset, dataloader, learner, etc. The new API will internally create everything `tsai` needs to be able to train a model.\n",
    "\n",
    "Note: these APIs don't replace the previous ones (get_ts_dls, ts_learner, etc) which will continue to work as always. They just offer a simpler API that is easier to use. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from tsai.imports import *\n",
    "from tsai.learner import *\n",
    "from tsai.data.all import *\n",
    "from tsai.models.InceptionTimePlus import *\n",
    "from tsai.models.utils import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TSClassifier API\n",
    "***\n",
    "\n",
    "**Commonly used arguments:**\n",
    "    \n",
    "* **X:** array-like of shape (n_samples, n_steps) or (n_samples, n_features, n_steps) with the input time series samples. Internally, they will be converted to torch tensors.\n",
    "* **y:** array-like of shape (n_samples), (n_samples, n_outputs) or (n_samples, n_features, n_outputs) with the target. Internally, they will be converted to torch tensors. Default=None. None is used for unlabeled datasets. \n",
    "* **splits:** lists of indices used to split data between train and validation. Default=None. If no splits are passed, data will be split 80:20 between train and test without shuffling.\n",
    "* **tfms:** item transforms that will be applied to each sample individually. Default:`[None, TSClassification()]` which is commonly used in most single label datasets. \n",
    "* **batch_tfms:** transforms applied to each batch. Default=None. \n",
    "* **bs:** batch size (if batch_size is provided then batch_size will override bs). An int or a list of ints can be passed. Default=`[64, 128]`. If a list of ints, the first one will be used for training, and the second for the valid (batch size can be larger as it doesn't require backpropagation which consumes more memory). \n",
    "* **arch:** indicates which architecture will be used. Default: InceptionTimePlus.\n",
    "* **arch_config:** keyword arguments passed to the selected architecture. Default={}.\n",
    "* **pretrained:** indicates if pretrained model weights will be used. Default=False.\n",
    "* **weights_path:** indicates the path to the pretrained weights in case they are used.\n",
    "* **loss_func:** allows you to pass any loss function. Default=None (in which case CrossEntropyLossFlat() is applied).\n",
    "* **opt_func:** allows you to pass an optimizer. Default=Adam.\n",
    "* **lr:** learning rate. Default=0.001.\n",
    "* **metrics:** list of metrics passed to the Learner. Default=accuracy.\n",
    "* **cbs:** list of callbacks passed to the Learner. Default=None.\n",
    "* **wd:** is the default weight decay used when training the model. Default=None.\n",
    "\n",
    "**Less frequently used arguments:**\n",
    "\n",
    "* **sel_vars:** used to select which of the features in multivariate datasets are used. Default=None means all features are used. If necessary a list-like of indices can be used (eg.`[0,3,5]`).\n",
    "* **sel_steps:** used to select the steps used. Default=None means all steps are used. If necessary a list-like of indices can be used (eg. `slice(-50, None)` will select the last 50 steps from each time series).\n",
    "* **weights:** indicates a sample weight per instance. Used to pass pass a probability to the train dataloader sampler. Samples with more weight will be selected more often during training. \n",
    "* **partial_n:** select randomly partial quantity of data at each epoch. Used to reduce the training size (for example for testing purposes). int or float can be used.\n",
    "* **inplace:** indicates whether tfms are applied during instantiation or on-the-fly. Default=True, which means that tfms will be applied during instantiation. This results in a faster training, but it can only be used when data fits in memory. Otherwise set it to False. \n",
    "* **shuffle_train:** indicates whether to shuffle the training set every time the dataloader is fully read/iterated or not. This doesn't have an impact on the validation set which is never shuffled. Default=True.\n",
    "* **drop_last:** if True the last incomplete training batch is dropped (thus ensuring training batches of equal size). This doesn't have an impact on the validation set where samples are never dropped. Default=True.\n",
    "* **num_workers:** num_workers (int): how many subprocesses to use for data loading. 0 means that the data will be loaded in the main process. Default=0. \n",
    "* **do_setup:** ndicates if the Pipeline.setup method should be called during initialization. Default=True.\n",
    "* **device:** Defaults to default_device() which is CUDA by default. You can specify device as `torch.device('cpu').\n",
    "* **verbose:** controls the verbosity when fitting and predicting.\n",
    "* **exclude_head:** indicates whether the head of the pretrained model needs to be removed or not. Default=True.\n",
    "* **cut:** indicates the position where the pretrained model head needs to be cut. Defaults=-1.\n",
    "* **init:** allows you to set to None (no initialization applied), set to True (in which case nn.init.kaiming_normal_ will be applied) or pass an initialization. Default=None.\n",
    "* **splitter:** To do transfer learning, you need to pass a splitter to Learner. This should be a function taking the model and returning a collection of parameter groups, e.g. a list of list of parameters. Default=trainable_params. If the model has a backbone and a head, it will then be split in those 2 groups.\n",
    "* **path** and **model_dir:** are used to save and/or load models. Often path will be inferred from dls, but you can override it or pass a Path object to model_dir.\n",
    "* **wd_bn_bias:** controls if weight decay is applied to BatchNorm layers and bias. Default=False.\n",
    "train_bn=True\n",
    "* **moms:** the default momentums used in Learner.fit_one_cycle. Default=(0.95, 0.85, 0.95)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\n",
    "class TSClassifier(Learner):\n",
    "    def __init__(self, X, y=None, splits=None, tfms=None, inplace=True, sel_vars=None, sel_steps=None, weights=None, partial_n=None, \n",
    "                 bs=[64, 128], batch_size=None, batch_tfms=None, shuffle_train=True, drop_last=True, num_workers=0, do_setup=True, device=None,\n",
    "                 arch=None, arch_config={}, pretrained=False, weights_path=None, exclude_head=True, cut=-1, init=None,\n",
    "                 loss_func=None, opt_func=Adam, lr=0.001, metrics=accuracy, cbs=None, wd=None, wd_bn_bias=False,\n",
    "                 train_bn=True, moms=(0.95, 0.85, 0.95),  path='.', model_dir='models', splitter=trainable_params, verbose=False):\n",
    "\n",
    "        #Splits\n",
    "        if splits is None: splits = TSSplitter()(X)\n",
    "            \n",
    "        # Item tfms\n",
    "        if tfms is None: tfms = [None, TSClassification()]\n",
    "\n",
    "        # Batch size\n",
    "        if batch_size is not None:\n",
    "            bs = batch_size\n",
    "\n",
    "        # DataLoaders\n",
    "        dls = get_ts_dls(X, y=y, splits=splits, sel_vars=sel_vars, sel_steps=sel_steps, tfms=tfms, inplace=inplace, \n",
    "                         path=path, bs=bs, batch_tfms=batch_tfms, num_workers=num_workers, weights=weights, partial_n=partial_n, \n",
    "                         device=device, shuffle_train=shuffle_train, drop_last=drop_last)\n",
    "        \n",
    "        if loss_func is None:\n",
    "            if hasattr(dls, 'loss_func'): loss_func = dls.loss_func\n",
    "            elif hasattr(dls, 'cat') and not dls.cat: loss_func = MSELossFlat()\n",
    "            elif hasattr(dls, 'train_ds') and hasattr(dls.train_ds, 'loss_func'): loss_func = dls.train_ds.loss_func\n",
    "            else: loss_func = CrossEntropyLossFlat()\n",
    "        \n",
    "        # Model\n",
    "        if init is True:\n",
    "            init = nn.init.kaiming_normal_\n",
    "        if arch is None:\n",
    "            arch = InceptionTimePlus\n",
    "        elif isinstance(arch, str): arch = get_arch(arch)\n",
    "        if 'xresnet' in arch.__name__.lower() and not '1d' in arch.__name__.lower():\n",
    "            model = build_tsimage_model(arch, dls=dls, pretrained=pretrained, init=init, device=device, verbose=verbose, arch_config=arch_config)\n",
    "        elif 'tabularmodel' in arch.__name__.lower():\n",
    "            build_tabular_model(arch, dls=dls, device=device, arch_config=arch_config)\n",
    "        else:\n",
    "            model = build_ts_model(arch, dls=dls, device=device, verbose=verbose, pretrained=pretrained, weights_path=weights_path,\n",
    "                                   exclude_head=exclude_head, cut=cut, init=init, arch_config=arch_config)\n",
    "        setattr(model, \"__name__\", arch.__name__)\n",
    "        try:\n",
    "            model[0], model[1]\n",
    "            splitter = ts_splitter\n",
    "        except:\n",
    "            pass\n",
    "\n",
    "        super().__init__(dls, model, loss_func=loss_func, opt_func=opt_func, lr=lr, cbs=cbs, metrics=metrics, path=path, splitter=splitter,\n",
    "                         model_dir=model_dir, wd=wd, wd_bn_bias=wd_bn_bias, train_bn=train_bn, moms=moms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>1.363282</td>\n",
       "      <td>1.391735</td>\n",
       "      <td>0.133333</td>\n",
       "      <td>00:09</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from tsai.models.InceptionTimePlus import *\n",
    "X, y, splits = get_classification_data('OliveOil', split_data=False)\n",
    "batch_tfms = [TSStandardize(by_sample=True)]\n",
    "learn = TSClassifier(X, y, splits=splits, batch_tfms=batch_tfms, metrics=accuracy, arch=InceptionTimePlus, arch_config=dict(fc_dropout=.5))\n",
    "learn.fit_one_cycle(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TSRegressor API\n",
    "***\n",
    "\n",
    "**Commonly used arguments:**\n",
    "    \n",
    "* **X:** array-like of shape (n_samples, n_steps) or (n_samples, n_features, n_steps) with the input time series samples. Internally, they will be converted to torch tensors.\n",
    "* **y:** array-like of shape (n_samples), (n_samples, n_outputs) or (n_samples, n_features, n_outputs) with the target. Internally, they will be converted to torch tensors. Default=None. None is used for unlabeled datasets. \n",
    "* **splits:** lists of indices used to split data between train and validation. Default=None. If no splits are passed, data will be split 80:20 between train and test without shuffling.\n",
    "* **tfms:** item transforms that will be applied to each sample individually. Default=`[None, TSRegression()]` which is commonly used in most single label datasets. \n",
    "* **batch_tfms:** transforms applied to each batch. Default=None. \n",
    "* **bs:** batch size (if batch_size is provided then batch_size will override bs). An int or a list of ints can be passed. Default=`[64, 128]`. If a list of ints, the first one will be used for training, and the second for the valid (batch size can be larger as it doesn't require backpropagation which consumes more memory). \n",
    "* **arch:** indicates which architecture will be used. Default: InceptionTimePlus.\n",
    "* **arch_config:** keyword arguments passed to the selected architecture. Default={}.\n",
    "* **pretrained:** indicates if pretrained model weights will be used. Default=False.\n",
    "* **weights_path:** indicates the path to the pretrained weights in case they are used.\n",
    "* **loss_func:** allows you to pass any loss function. Default=None (in which case CrossEntropyLossFlat() is applied).\n",
    "* **opt_func:** allows you to pass an optimizer. Default=Adam.\n",
    "* **lr:** learning rate. Default=0.001.\n",
    "* **metrics:** list of metrics passed to the Learner. Default=None.\n",
    "* **cbs:** list of callbacks passed to the Learner. Default=None.\n",
    "* **wd:** is the default weight decay used when training the model. Default=None.\n",
    "\n",
    "**Less frequently used arguments:**\n",
    "\n",
    "* **sel_vars:** used to select which of the features in multivariate datasets are used. Default=None means all features are used. If necessary a list-like of indices can be used (eg.`[0,3,5]`).\n",
    "* **sel_steps:** used to select the steps used. Default=None means all steps are used. If necessary a list-like of indices can be used (eg. `slice(-50, None)` will select the last 50 steps from each time series).\n",
    "* **weights:** indicates a sample weight per instance. Used to pass pass a probability to the train dataloader sampler. Samples with more weight will be selected more often during training. \n",
    "* **partial_n:** select randomly partial quantity of data at each epoch. Used to reduce the training size (for example for testing purposes). int or float can be used.\n",
    "* **inplace:** indicates whether tfms are applied during instantiation or on-the-fly. Default=True, which means that tfms will be applied during instantiation. This results in a faster training, but it can only be used when data fits in memory. Otherwise set it to False. \n",
    "* **shuffle_train:** indicates whether to shuffle the training set every time the dataloader is fully read/iterated or not. This doesn't have an impact on the validation set which is never shuffled. Default=True.\n",
    "* **drop_last:** if True the last incomplete training batch is dropped (thus ensuring training batches of equal size). This doesn't have an impact on the validation set where samples are never dropped. Default=True.\n",
    "* **num_workers:** num_workers (int): how many subprocesses to use for data loading. 0 means that the data will be loaded in the main process. Default=0. \n",
    "* **do_setup:** ndicates if the Pipeline.setup method should be called during initialization. Default=True.\n",
    "* **device:** Defaults to default_device() which is CUDA by default. You can specify device as `torch.device('cpu').\n",
    "* **verbose:** controls the verbosity when fitting and predicting.\n",
    "* **exclude_head:** indicates whether the head of the pretrained model needs to be removed or not. Default=True.\n",
    "* **cut:** indicates the position where the pretrained model head needs to be cut. Defaults=-1.\n",
    "* **init:** allows you to set to None (no initialization applied), set to True (in which case nn.init.kaiming_normal_ will be applied) or pass an initialization. Default=None.\n",
    "* **splitter:** To do transfer learning, you need to pass a splitter to Learner. This should be a function taking the model and returning a collection of parameter groups, e.g. a list of list of parameters. Default=trainable_params. If the model has a backbone and a head, it will then be split in those 2 groups.\n",
    "* **path** and **model_dir:** are used to save and/or load models. Often path will be inferred from dls, but you can override it or pass a Path object to model_dir.\n",
    "* **wd_bn_bias:** controls if weight decay is applied to BatchNorm layers and bias. Default=False.\n",
    "train_bn=True\n",
    "* **moms:** the default momentums used in Learner.fit_one_cycle. Default=(0.95, 0.85, 0.95)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    " \n",
    "        \n",
    "class TSRegressor(Learner):\n",
    "    def __init__(self, X, y=None, splits=None, tfms=None, inplace=True, sel_vars=None, sel_steps=None, weights=None, partial_n=None, \n",
    "                 bs=[64, 128], batch_size=None, batch_tfms=None, shuffle_train=True, drop_last=True, num_workers=0, do_setup=True, device=None,\n",
    "                 arch=None, arch_config={}, pretrained=False, weights_path=None, exclude_head=True, cut=-1, init=None,\n",
    "                 loss_func=None, opt_func=Adam, lr=0.001, metrics=None, cbs=None, wd=None, wd_bn_bias=False,\n",
    "                 train_bn=True, moms=(0.95, 0.85, 0.95),  path='.', model_dir='models', splitter=trainable_params, verbose=False):\n",
    "\n",
    "        #Splits\n",
    "        if splits is None: splits = TSSplitter()(X)\n",
    "            \n",
    "        # Item tfms\n",
    "        if tfms is None: tfms = [None, TSRegression()]\n",
    "\n",
    "        # Batch size\n",
    "        if batch_size is not None:\n",
    "            bs = batch_size\n",
    "\n",
    "        # DataLoaders\n",
    "        dls = get_ts_dls(X, y=y, splits=splits, sel_vars=sel_vars, sel_steps=sel_steps, tfms=tfms, inplace=inplace, \n",
    "                         path=path, bs=bs, batch_tfms=batch_tfms, num_workers=num_workers, weights=weights, partial_n=partial_n, \n",
    "                         device=device, shuffle_train=shuffle_train, drop_last=drop_last)\n",
    "\n",
    "        if loss_func is None:\n",
    "            if hasattr(dls, 'loss_func'): loss_func = dls.loss_func\n",
    "            elif hasattr(dls, 'cat') and not dls.cat: loss_func = MSELossFlat()\n",
    "            elif hasattr(dls, 'train_ds') and hasattr(dls.train_ds, 'loss_func'): loss_func = dls.train_ds.loss_func\n",
    "            else: loss_func = MSELossFlat()\n",
    "                \n",
    "        # Model\n",
    "        if init is True:\n",
    "            init = nn.init.kaiming_normal_\n",
    "        if arch is None:\n",
    "            arch = InceptionTimePlus\n",
    "        elif isinstance(arch, str): arch = get_arch(arch)\n",
    "        if 'xresnet' in arch.__name__.lower() and not '1d' in arch.__name__.lower():\n",
    "            model = build_tsimage_model(arch, dls=dls, pretrained=pretrained, init=init, device=device, verbose=verbose, arch_config=arch_config)\n",
    "        elif 'tabularmodel' in arch.__name__.lower():\n",
    "            build_tabular_model(arch, dls=dls, device=device, arch_config=arch_config)\n",
    "        else:\n",
    "            model = build_ts_model(arch, dls=dls, device=device, verbose=verbose, pretrained=pretrained, weights_path=weights_path,\n",
    "                               exclude_head=exclude_head, cut=cut, init=init, arch_config=arch_config)\n",
    "        setattr(model, \"__name__\", arch.__name__)\n",
    "        try:\n",
    "            model[0], model[1]\n",
    "            splitter = ts_splitter\n",
    "        except:\n",
    "            pass\n",
    "\n",
    "        super().__init__(dls, model, loss_func=loss_func, opt_func=opt_func, lr=lr, cbs=cbs, metrics=metrics, path=path, splitter=splitter,\n",
    "                         model_dir=model_dir, wd=wd, wd_bn_bias=wd_bn_bias, train_bn=train_bn, moms=moms)      "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>mae</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>225.968323</td>\n",
       "      <td>213.152908</td>\n",
       "      <td>14.189029</td>\n",
       "      <td>00:11</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from tsai.models.TST import *\n",
    "X, y, splits = get_regression_data('AppliancesEnergy', split_data=False)\n",
    "if X is not None: # This is to prevent a test fail when the data server is not available\n",
    "    batch_tfms = [TSStandardize()]\n",
    "    learn = TSRegressor(X, y, splits=splits, batch_tfms=batch_tfms, arch=TST, metrics=mae, bs=512)\n",
    "    learn.fit_one_cycle(1, 1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TSForecaster API\n",
    "***\n",
    "**Commonly used arguments:**\n",
    "    \n",
    "* **X:** array-like of shape (n_samples, n_steps) or (n_samples, n_features, n_steps) with the input time series samples. Internally, they will be converted to torch tensors.\n",
    "* **y:** array-like of shape (n_samples), (n_samples, n_outputs) or (n_samples, n_features, n_outputs) with the target. Internally, they will be converted to torch tensors. Default=None. None is used for unlabeled datasets. \n",
    "* **splits:** lists of indices used to split data between train and validation. Default=None. If no splits are passed, data will be split 80:20 between train and test without shuffling.\n",
    "* **tfms:** item transforms that will be applied to each sample individually. Default=`[None, TSForecasting()]` which is commonly used in most single label datasets. \n",
    "* **batch_tfms:** transforms applied to each batch. Default=None. \n",
    "* **bs:** batch size (if batch_size is provided then batch_size will override bs). An int or a list of ints can be passed. Default=`[64, 128]`. If a list of ints, the first one will be used for training, and the second for the valid (batch size can be larger as it doesn't require backpropagation which consumes more memory). \n",
    "* **arch:** indicates which architecture will be used. Default: InceptionTimePlus.\n",
    "* **arch_config:** keyword arguments passed to the selected architecture. Default={}.\n",
    "* **pretrained:** indicates if pretrained model weights will be used. Default=False.\n",
    "* **weights_path:** indicates the path to the pretrained weights in case they are used.\n",
    "* **loss_func:** allows you to pass any loss function. Default=None (in which case CrossEntropyLossFlat() is applied).\n",
    "* **opt_func:** allows you to pass an optimizer. Default=Adam.\n",
    "* **lr:** learning rate. Default=0.001.\n",
    "* **metrics:** list of metrics passed to the Learner. Default=None.\n",
    "* **cbs:** list of callbacks passed to the Learner. Default=None.\n",
    "* **wd:** is the default weight decay used when training the model. Default=None.\n",
    "\n",
    "**Less frequently used arguments:**\n",
    "\n",
    "* **sel_vars:** used to select which of the features in multivariate datasets are used. Default=None means all features are used. If necessary a list-like of indices can be used (eg.`[0,3,5]`).\n",
    "* **sel_steps:** used to select the steps used. Default=None means all steps are used. If necessary a list-like of indices can be used (eg. `slice(-50, None)` will select the last 50 steps from each time series).\n",
    "* **weights:** indicates a sample weight per instance. Used to pass pass a probability to the train dataloader sampler. Samples with more weight will be selected more often during training. \n",
    "* **partial_n:** select randomly partial quantity of data at each epoch. Used to reduce the training size (for example for testing purposes). int or float can be used.\n",
    "* **inplace:** indicates whether tfms are applied during instantiation or on-the-fly. Default=True, which means that tfms will be applied during instantiation. This results in a faster training, but it can only be used when data fits in memory. Otherwise set it to False. \n",
    "* **shuffle_train:** indicates whether to shuffle the training set every time the dataloader is fully read/iterated or not. This doesn't have an impact on the validation set which is never shuffled. Default=True.\n",
    "* **drop_last:** if True the last incomplete training batch is dropped (thus ensuring training batches of equal size). This doesn't have an impact on the validation set where samples are never dropped. Default=True.\n",
    "* **num_workers:** num_workers (int): how many subprocesses to use for data loading. 0 means that the data will be loaded in the main process. Default=None. \n",
    "* **do_setup:** ndicates if the Pipeline.setup method should be called during initialization. Default=True.\n",
    "* **device:** Defaults to default_device() which is CUDA by default. You can specify device as `torch.device('cpu').\n",
    "* **verbose:** controls the verbosity when fitting and predicting.\n",
    "* **exclude_head:** indicates whether the head of the pretrained model needs to be removed or not. Default=True.\n",
    "* **cut:** indicates the position where the pretrained model head needs to be cut. Defaults=-1.\n",
    "* **init:** allows you to set to None (no initialization applied), set to True (in which case nn.init.kaiming_normal_ will be applied) or pass an initialization. Default=None.\n",
    "* **splitter:** To do transfer learning, you need to pass a splitter to Learner. This should be a function taking the model and returning a collection of parameter groups, e.g. a list of list of parameters. Default=trainable_params. If the model has a backbone and a head, it will then be split in those 2 groups.\n",
    "* **path** and **model_dir:** are used to save and/or load models. Often path will be inferred from dls, but you can override it or pass a Path object to model_dir.\n",
    "* **wd_bn_bias:** controls if weight decay is applied to BatchNorm layers and bias. Default=False.\n",
    "train_bn=True\n",
    "* **moms:** the default momentums used in Learner.fit_one_cycle. Default=(0.95, 0.85, 0.95)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export  \n",
    "        \n",
    "class TSForecaster(Learner):\n",
    "    def __init__(self, X, y=None, splits=None, tfms=None, inplace=True, sel_vars=None, sel_steps=None, weights=None, partial_n=None, \n",
    "                 bs=[64, 128], batch_size=None, batch_tfms=None, shuffle_train=True, drop_last=True, num_workers=0, do_setup=True, device=None,\n",
    "                 arch=None, arch_config={}, pretrained=False, weights_path=None, exclude_head=True, cut=-1, init=None,\n",
    "                 loss_func=None, opt_func=Adam, lr=0.001, metrics=None, cbs=None, wd=None, wd_bn_bias=False,\n",
    "                 train_bn=True, moms=(0.95, 0.85, 0.95),  path='.', model_dir='models', splitter=trainable_params, verbose=False):\n",
    "\n",
    "        #Splits\n",
    "        if splits is None: splits = TSSplitter()(X)\n",
    "            \n",
    "        # Item tfms\n",
    "        if tfms is None: tfms = [None, TSForecasting()]\n",
    "\n",
    "        # Batch size\n",
    "        if batch_size is not None:\n",
    "            bs = batch_size\n",
    "\n",
    "        # DataLoaders\n",
    "        dls = get_ts_dls(X, y=y, splits=splits, sel_vars=sel_vars, sel_steps=sel_steps, tfms=tfms, inplace=inplace, \n",
    "                         path=path, bs=bs, batch_tfms=batch_tfms, num_workers=num_workers, weights=weights, partial_n=partial_n, \n",
    "                         device=device, shuffle_train=shuffle_train, drop_last=drop_last)\n",
    "        \n",
    "        if loss_func is None:\n",
    "            if hasattr(dls, 'loss_func'): loss_func = dls.loss_func\n",
    "            elif hasattr(dls, 'cat') and not dls.cat: loss_func = MSELossFlat()\n",
    "            elif hasattr(dls, 'train_ds') and hasattr(dls.train_ds, 'loss_func'): loss_func = dls.train_ds.loss_func\n",
    "            else: loss_func = MSELossFlat()\n",
    "        \n",
    "        # Model\n",
    "        if init is True:\n",
    "            init = nn.init.kaiming_normal_\n",
    "        if arch is None:\n",
    "            arch = InceptionTimePlus\n",
    "        elif isinstance(arch, str): arch = get_arch(arch)\n",
    "        if 'xresnet' in arch.__name__.lower() and not '1d' in arch.__name__.lower():\n",
    "            model = build_tsimage_model(arch, dls=dls, pretrained=pretrained, init=init, device=device, verbose=verbose, arch_config=arch_config)\n",
    "        elif 'tabularmodel' in arch.__name__.lower():\n",
    "            build_tabular_model(arch, dls=dls, device=device, arch_config=arch_config)\n",
    "        else:\n",
    "            model = build_ts_model(arch, dls=dls, device=device, verbose=verbose, pretrained=pretrained, weights_path=weights_path,\n",
    "                               exclude_head=exclude_head, cut=cut, init=init, arch_config=arch_config)\n",
    "        setattr(model, \"__name__\", arch.__name__)\n",
    "        try:\n",
    "            model[0], model[1]\n",
    "            splitter = ts_splitter\n",
    "        except:\n",
    "            pass\n",
    "\n",
    "        super().__init__(dls, model, loss_func=loss_func, opt_func=opt_func, lr=lr, cbs=cbs, metrics=metrics, path=path, splitter=splitter,\n",
    "                         model_dir=model_dir, wd=wd, wd_bn_bias=wd_bn_bias, train_bn=train_bn, moms=moms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset: Sunspots\n",
      "downloading data...\n",
      "...data downloaded. Path = data/forecasting/Sunspots.csv\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABAYAAABKCAYAAAAoj1bdAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAAAsTAAALEwEAmpwYAAAQiklEQVR4nO3df5DU9X3H8dfr7pAfxaAnJwinwYLRw4t35AxoShBJi8bGpKj4C38kkZhxkmkaM5G0zagBMrWdtnFak4zREGhNiI6oVYdWaUTwxwQV6kVFEGLJgOG8I3eISNBb7t0/9nvpZt29uyV7t9zt8zFzw3e/389+vu8v957P3r738/2sI0IAAAAAAKA8VZQ6AAAAAAAAUDoUBgAAAAAAKGMUBgAAAAAAKGMUBgAAAAAAKGMUBgAAAAAAKGMUBgAAAAAAKGMUBgAAg5btJ20vTLYX2H78D+hrku2wXZU8/k/b1xYpzo/b3prxeIftPy1G30l/r9ieXaz+AABAeaEwAAAoKdszbT9r+y3b7bafsf3RQvuJiB9HxNyMfsP2lMONKyI+GREremvXl/NExFMRcerhxpJ1vuW2l2b1f3pEPFmM/gEAQPmpKnUAAIDyZfsDkh6VdIOk+yQdJenjkt4tZVzFZLsqIlKljgMAACAfZgwAAErpQ5IUESsj4lBE/DYiHo+IX0iS7c8mMwjuSGYUbLH9iVwdJW2fTrbXJ7ubbe+3fVmO9pW2/9H2HtuvS/rzrOOZtylMsb0uiWGP7Xvzncf2bNu7bC+y3SLpR937skL4qO3Ntjts/8j2iOzryIglkhiul7RA0k3J+R5Jjv/u1gTbw23fbvvXyc/ttocnx7pj+5rtVtu7bX+u198SAAAY0igMAABK6TVJh2yvsP1J28fmaDND0i8ljZV0i6QHbFf31GlEzEo2GyJidETcm6PZFyR9StI0SWdKuqSHLpdIelzSsZJqJf1rL+cZL6la0gclXZ+nzwWSzpM0WekCyTd7uqbkfD+Q9GNJ/5Cc78Iczf5W0lmSGiU1SJqe1fd4SWMkTZR0naTv5vl/BwAAZYLCAACgZCJin6SZkkLSXZLabD9se1xGs1ZJt0dEZ/LGe6uyPt0/TJcm/e6MiHZJf9dD206l3+RPiIiDEfF0D20lqUvSLRHxbkT8Nk+bOzLO/W1JVxR6AXkskLQ4Ilojok3StyRdnXG8MzneGRGrJe2XVJT1DwAAwOBEYQAAUFIR8WpEfDYiaiXVS5og6faMJm9ERGQ8/lXS5g81QdLOrH7zuUmSJT2XfAPA53vpuy0iDvbSJvvcxbgmJf1kXkt237/JWvPggKTRRTo3AAAYhCgMAACOGBGxRdJypQsE3SbadsbjkyT9ugin2y3pxKx+88XVEhFfiIgJkr4o6Xu9fBNB9HCsW/a5u6/pHUmjug/YHl9g379WenZDrr4BAADeh8IAAKBkbJ+WLIRXmzw+Uekp9T/PaHa8pL+0Pcz2fEl1klb3ofs3Jf1xD8fvS/qtTe6x/0YPcc7vjlFSh9Jvzrv6eJ58vpScu1rpdQG61ydolnS67cZkQcJbs57X2/lWSvqm7RrbYyXdLOmew4gPAACUCQoDAIBSelvpxQU32H5H6YLAy5K+ltFmg6RTJO1R+l78SyLiN33o+1ZJK2zvtX1pjuN3SXpM6TfimyQ90ENfH01i3C/pYUlfiYjX+3iefH6i9IKGryu9uOJSSYqI1yQtlvTfkrZJyl7P4IeSpibneyhHv0slvSDpF5JeSq5taQFxAQCAMuPfv20TAIAjh+3PSloYETNLHQsAAMBQxYwBAAAAAADKGIUBAAAAAADKGLcSAAAAAABQxpgxAAAAAABAGaMwAAAAAABAGavqj07tsSFN6o+uAQAAAKAkRtW9WuoQjggHXj2wJyJqSh0HiqdfCgPposAL/dM1AAAAAJTAafc0lTqEI8Kmpk2/KnUMKC5uJQAAAAAAoIxRGAAAAAAAoIxRGAAAAAAAoIz10xoDAAAAAAAcuTZu3Hh8VVXV3ZLqNbQ/NO+S9HIqlVrY1NTUmqsBhQEAAAAAQNmpqqq6e/z48XU1NTUdFRUVUep4+ktXV5fb2tqmtrS03C3p07naDOWqCAAAAAAA+dTX1NTsG8pFAUmqqKiImpqat5SeGZG7zQDGAwAAAADAkaJiqBcFuiXXmff9P7cSAAAAAAAwwFpaWipnz559qiTt2bNnWEVFRVRXV6ck6cUXX3x1xIgReYsW69evH7Vs2bLjli9fvrMYsfRaGLC9TNKnJLVGRN6pBwAAAAAADFa2morZX4Q29nR8/Pjxh7Zs2bJZkm688cYJo0ePPrR48eI3u493dnZq2LBhOZ87a9asA7NmzTpQrFj7civBcknnF+uEAAAAAADg/S6++OJJV1555UlnnHHGaTfccEPt2rVrRzU2Np5WV1c3ddq0aac1NzcPl6RHH3306HPPPXeKlC4qzJ8/f9L06dNPra2t/fDSpUuPL/S8vc4YiIj1ticVfEUAAAAAAKAgu3fvPmrTpk1bqqqq1N7eXvH8889vGTZsmB566KGjb7rpptrHHnvsl9nP2b59+4hnn3126969eyvr6urqv/71r7cNHz68z+snFG2NAdvXS7o+/eikYnULAAAAAEDZuOiiizqqqtJv1dvb2ysvu+yyk3fs2DHCdnR2djrXc+bOnbt35MiRMXLkyFR1dXXnrl27qiZPntzZ13MW7VsJIuIHEXFmRJwp1RSrWwAAAAAAysbo0aO7urcXLVo08Zxzznl727ZtrzzyyCPb33vvvZzv4TNnB1RWViqVSuUsIOTD1xUCAAAAAHAE2rdvX2Vtbe17knTnnXeO7a/zUBgAAAAAAOAItGjRopZbb721tq6ubmoqleq38zii5/UIbK+UNFvSWElvSrolIn7Y83PODOmFYsUIAAAAACX3kY1F/Ta7QWtT06aN6VvIB7fm5uYdDQ0Ne0odx0Bpbm4e29DQMCnXsb58K8EVRY8IAAAAAAAcEbiVAAAAAACAMkZhAAAAAACAMkZhAAAAAACAMkZhAAAAAACAMkZhAAAAAACAMkZhAAAAAACAATZjxowPrVq16gOZ+xYvXnz8ggULTsrVfvr06aeuX79+lCSdc845U/bs2VOZ3ebGG2+ccPPNN48rNJZev64QAAAAAIChrmlTU1Mx+9v4kY0bezo+f/789pUrV1ZffPHF+7r3rVq1qvq2227b1Vvf69at216MGLsxYwAAAAAAgAF29dVXdzzxxBNjDh48aEnaunXrUa2trcPuueee6vr6+ropU6ac/tWvfnVCrudOnDjxw7t3766SpEWLFo2fNGlSfVNT06nbtm0bfjix9NOMgY37JW/tn74xBI2VtKfUQWBQIWdQCPIFhSJnUChypkxsKs7nyUMhXz5Y6gCGgnHjxh1qaGh45/777x9z1VVX7V2xYkX1hRde2LFkyZLd48aNO5RKpfSxj33s1A0bNoycMWPGb3P18dRTT4168MEHq1966aXNnZ2damxsnDpt2rQDhcbSX7cSbI2IM/upbwwxtl8gX1AIcgaFIF9QKHIGhSJnUAjyBZkuvfTS9nvvvffYq666au8DDzxQfdddd+1YsWJF9fLly8emUim3tbUNa25uHpGvMLB27drRF1xwwd6jjz66S5Lmzp2793Di4FYCAAAAAABK4Morr9z7zDPPfODpp58edfDgwYqamprUHXfcMW7dunWvvfbaa5vnzJnz1sGDB/v9fTuFAQAAAAAASmDMmDFdZ5999tsLFy6cNG/evPaOjo7KkSNHdlVXVx/auXNn1ZNPPjmmp+fPmTNn/+rVq4/Zv3+/Ozo6KtasWXPM4cTRX7cS/KCf+sXQRL6gUOQMCkG+oFDkDApFzqAQ5At+z+WXX95+zTXXTF65cuXr06ZNO1hfX39g8uTJ9SeccMJ7TU1N+3t67syZMw/Mmzevvb6+/vTjjjuu84wzznjncGJwRBxe9AAAAAAADFLNzc07GhoaBvtCkH3W3Nw8tqGhYVKuY9xKAAAAAABAGStqYcD2+ba32t5u+xvF7BuDm+0dtl+y/aLtF5J91bbX2N6W/Htsst+2/yXJo1/Y/khpo0d/s73MdqvtlzP2FZwftq9N2m+zfW0prgUDI0/O3Gr7jWScedH2BRnH/jrJma22z8vYz+tWGbB9ou21tjfbfsX2V5L9jDPIqYecYZzB+9geYfs5281Jvnwr2X+y7Q3J7/5e20cl+4cnj7cnxydl9JUzj4D+VrTCgO1KSd+V9ElJUyVdYXtqsfrHkHBuRDRmfD3LNyT9LCJOkfSz5LGUzqFTkp/rJX1/wCPFQFsu6fysfQXlh+1qSbdImiFpuqRbuv/Ix5C0XO/PGUn6TjLONEbEaklKXosul3R68pzv2a7kdauspCR9LSKmSjpL0peS3zXjDPLJlzMS4wze711JcyKiQVKjpPNtnyXp75XOlymSOiRdl7S/TlJHsv87Sbu8eTSQF4LyVcwZA9MlbY+I1yPiPUk/lfSZIvaPoeczklYk2ysk/UXG/n+LtJ9LOsb2CSWIDwMkItZLas/aXWh+nCdpTUS0R0SHpDXK/cYRQ0CenMnnM5J+GhHvRsT/Stqu9GsWr1tlIiJ2R8SmZPttSa9KmijGGeTRQ87kwzhTxpKxonuBuGHJT0iaI+n+ZH/2GNM99twv6RO2rfx5hP7T1dXV5VIHMRCS6+zKd7yYhYGJknZmPN6lngdQlJeQ9LjtjbavT/aNi4jdyXaLpHHJNrkEqfD8IG8gSV9Opn4vy/gkl5zB7yRTdqdJ2iDGGfRBVs5IjDPIIZkh8qKkVqWLhr+UtDciUkmTzN/97/IiOf6WpONEvpTCy21tbWOGenGgq6vLbW1tYyS9nK9Nf31dIZBtZkS8Yft4SWtsb8k8GBFhm6/IQE7kB/ro+5KWKF2IXCLpnyR9vqQR4Yhie7SkVZL+KiL2pT+gS2OcQS45coZxBjlFxCFJjbaPkfSgpNNKGxH6IpVKLWxpabm7paWlXkN7Yf4uSS+nUqmF+RoUszDwhqQTMx7XJvsARcQbyb+tth9UelrUm7ZPiIjdyRTN1qQ5uQSp8Px4Q9LsrP1PDkCcOEJExJvd27bvkvRo8rCnMYWxpkzYHqb0G7wfR8QDyW7GGeSVK2cYZ9CbiNhre62ks5W+DakqmRWQ+bvvzpddtqskjZH0G/E38IBrampqlfTpUsdxJChmVeR5Sackq28epfTCGQ8XsX8MUrb/yPbR3duS5io9jeVhSd0rOl8r6T+S7YclXZOsCn2WpLcypnqifBSaH49Jmmv72GRq59xkH8pE1lok8/T/0+UelnR5sgr0yUovKPeceN0qG8m9uz+U9GpE/HPGIcYZ5JQvZxhnkIvtmmSmgGyPlPRnSq9LsVbSJUmz7DGme+y5RNITERHKn0dAvyvajIGISNn+stIvkJWSlkXEK8XqH4PaOEkPJlM2qyT9JCL+y/bzku6zfZ2kX0m6NGm/WtIFSi+4ckDS5wY+ZAwk2yuV/hRurO1dSq/6fZsKyI+IaLe9ROk/wiRpcUT0dXE6DDJ5cma27Ualp/jukPRFSYqIV2zfJ2mz0iuNfymZ8ilet8rGn0i6WtJLyT3AkvQ3YpxBfvly5grGGeRwgqQVyTcIVEi6LyIetb1Z0k9tL5X0P0oXm5T8+++2tyu9kO7lUs95BPQ3p4tTAAAAAACgHA3lBRYAAAAAAEAvKAwAAAAAAFDGKAwAAAAAAFDGKAwAAAAAAFDGKAwAAAAAAFDGKAwAAAAAAFDGKAwAAAAAAFDGKAwAAAAAAFDG/g/s/SiQrGI0mAAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 1152x36 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>mae</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>11146.689453</td>\n",
       "      <td>404.269287</td>\n",
       "      <td>17.182791</td>\n",
       "      <td>00:07</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from tsai.models.TSTPlus import *\n",
    "ts = get_forecasting_time_series('Sunspots')\n",
    "if ts is not None: # This is to prevent a test fail when the data server is not available\n",
    "    X, y = SlidingWindowSplitter(60, horizon=1)(ts)\n",
    "    splits = TSSplitter(235)(y)\n",
    "    batch_tfms = [TSStandardize(by_var=True)]\n",
    "    learn = TSForecaster(X, y, splits=splits, batch_tfms=batch_tfms, arch=TST, arch_config=dict(fc_dropout=.5), metrics=mae, bs=512, partial_n=.1)\n",
    "    learn.fit_one_cycle(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from tsai.imports import create_scripts\n",
    "from tsai.export import get_nb_name\n",
    "nb_name = get_nb_name()\n",
    "create_scripts(nb_name);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
